from pathlib import Path
from typing import List, Dict, Any, TypedDict
from dataclasses import dataclass
import argparse
import re
import sys
import os
import json
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage, SystemMessage
from typing_extensions import Literal
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from models.llm_models import GLOBAL_LLM
from config import cfg


PROJECT_ROOT = Path(__file__).parent.parent.resolve()


class TranslationFeedback(BaseModel):
    """翻译反馈数据模型"""
    grade: Literal["excellent", "good", "acceptable", "poor"] = Field(
        description="评估翻译质量等级：excellent(优秀)、good(良好)、acceptable(可接受)、poor(较差)",
    )
    feedback: str = Field(
        description="如果翻译质量不佳，请提供有关如何改进的详细反馈。",
    )


@dataclass
class TranslationAgentInput:
    """Translation Agent Input"""
    input_path: str
    output_path: str
    target_language: str
    use_batch: bool = True


@dataclass
class TranslationAgentOutput:
    """Translation Agent Output"""
    input_path: str
    output_path: str
    target_language: str
    total_subtitles: int
    translated_subtitles: int
    failed_translations: int


class TranslationAgentState(TypedDict):
    """Translation Agent State"""
    input_path: str
    output_path: str
    target_language: str
    use_batch: bool
    all_subtitles: List[Dict[str, Any]]  # 所有字幕数据
    batch_translations: List[str]  # 批量翻译结果
    evaluation_results: List[Dict[str, Any]]  # 评估结果
    failed_indices: List[int]  # 不合格的索引
    final_translations: List[str]  # 最终翻译结果
    result: TranslationAgentOutput


@dataclass
class TranslationAgentConfig:
    """Translation Agent Config"""
    agent_name: str
    max_llm_context: int = 32000


class TranslationAgent:
    """Translation Agent"""
    
    def __init__(self, config: TranslationAgentConfig):
        self.config = config
        self.max_llm_context = config.max_llm_context
        self.app = self._build_graph()
        self.batch_app = self._build_batch_graph()
        self.llm = GLOBAL_LLM

        self.workspace_dir = PROJECT_ROOT / "agents_workspace" / self.config.agent_name
        if not self.workspace_dir.exists():
            self.workspace_dir.mkdir(parents=True, exist_ok=True)

    def _build_graph(self) -> StateGraph:
        """Build the individual translation graph"""
        workflow = StateGraph(TranslationAgentState)
        
        # 添加节点
        workflow.add_node("load_subtitles", self._load_subtitles)
        workflow.add_node("llm_call_translator", self._llm_call_translator)
        workflow.add_node("llm_call_evaluator", self._llm_call_evaluator)
        workflow.add_node("generate_result", self._generate_result)
        
        # 设置边
        workflow.set_entry_point("load_subtitles")
        workflow.add_edge("load_subtitles", "llm_call_translator")
        workflow.add_edge("llm_call_translator", "llm_call_evaluator")
        workflow.add_conditional_edges(
            "llm_call_evaluator",
            self._route_translation,
            {  
                "Accepted": "generate_result",
                "Rejected + Feedback": "llm_call_translator",
            },
        )
        workflow.add_edge("generate_result", END)
        
        return workflow.compile()
    
    def _build_batch_graph(self) -> StateGraph:
        """Build the batch translation graph"""
        workflow = StateGraph(TranslationAgentState)
        
        # 添加节点
        workflow.add_node("load_subtitles", self._load_subtitles)
        workflow.add_node("batch_translator", self._batch_translator)
        workflow.add_node("batch_evaluator", self._batch_evaluator)
        workflow.add_node("individual_retranslator", self._individual_retranslator)
        workflow.add_node("generate_result", self._generate_result)
        
        # 设置边
        workflow.set_entry_point("load_subtitles")
        workflow.add_edge("load_subtitles", "batch_translator")
        workflow.add_edge("batch_translator", "batch_evaluator")
        workflow.add_conditional_edges(
            "batch_evaluator",
            self._route_batch_processing,
            {  
                "Completed": "generate_result",
                "Retranslate": "individual_retranslator",
            },
        )
        workflow.add_edge("individual_retranslator", "generate_result")
        workflow.add_edge("generate_result", END)
        
        return workflow.compile()
    
    async def _batch_translator(self, state: TranslationAgentState) -> TranslationAgentState:
        """批量翻译所有台词"""
        
        # 提取所有需要翻译的文本
        texts_to_translate = []
        for subtitle in state["all_subtitles"]:
            text = subtitle.get("text", "").strip()
            if text:
                texts_to_translate.append(text)
        
        if not texts_to_translate:
            state["batch_translations"] = []
            return state
        
        # 将所有文本合并为一个批次
        combined_text = "\n\n".join([f"{i+1}. {text}" for i, text in enumerate(texts_to_translate)])
        
        # 批量翻译
        prompt = f"""请将以下台词批量翻译成{state['target_language']}。

要求：
1. 保持原有的编号格式
2. 每行翻译对应一个编号
3. 只输出翻译结果，不要添加其他内容
4. 保持原文的语气和风格

原文：
{combined_text}

翻译结果："""
        
        messages = [{"role": "user", "content": prompt}]
        response = await self.llm.a_run(messages)
        translated_content = response.content.strip()
        
        # 解析批量翻译结果
        batch_translations = []
        lines = translated_content.split('\n')
        
        print(f"LLM返回的翻译内容：\n{translated_content}")
        print(f"分割后的行数：{len(lines)}")
        
        for i, line in enumerate(lines):
            line = line.strip()
            if not line:
                continue
                
            # 移除编号前缀，提取翻译文本
            if '. ' in line:
                translation = line.split('. ', 1)[1].strip()
            else:
                translation = line.strip()
                
            if translation:
                batch_translations.append(translation)
                print(f"第{i+1}行解析结果：{translation}")
        
        print(f"最终解析出 {len(batch_translations)} 条翻译")
        state["batch_translations"] = batch_translations
        return state


    async def _batch_evaluator(self, state: TranslationAgentState) -> TranslationAgentState:
        """批量评估翻译质量"""
        
        # 准备批量评估的输入
        evaluation_pairs = []
        for i, (original, translated) in enumerate(zip(state["all_subtitles"], state["batch_translations"])):
            original_text = original.get("text", "").strip()
            if original_text and translated:
                evaluation_pairs.append(f"{i+1}. 原文：{original_text}\n   译文：{translated}")
        
        if not evaluation_pairs:
            state["evaluation_results"] = []
            state["failed_indices"] = []
            return state
        
        # 批量评估所有翻译
        combined_evaluation = "\n\n".join(evaluation_pairs)
        
        batch_evaluation_prompt = f"""请对以下批量翻译结果进行评估，目标语言是{state['target_language']}。

要求：
1. 对每个翻译给出质量等级：excellent(优秀)、good(良好)、acceptable(可接受)、poor(较差)
2. 如果翻译质量不佳，请提供改进建议
3. 请按以下格式输出评估结果：
编号: 等级 - 反馈

翻译内容：
{combined_evaluation}

评估结果："""

        messages = [{"role": "user", "content": batch_evaluation_prompt}]
        evaluation_response = await self.llm.a_run(messages)
        evaluation_text = evaluation_response.content.strip()
        
        # 解析批量评估结果
        evaluation_results = []
        failed_indices = []
        
        lines = evaluation_text.split('\n')
        for line in lines:
            line = line.strip()
            if not line or ':' not in line:
                continue
                
            try:
                # 解析格式：编号: 等级 - 反馈
                parts = line.split(':', 1)
                if len(parts) < 2:
                    continue
                    
                index_part = parts[0].strip()
                evaluation_part = parts[1].strip()
                
                # 提取编号
                index = int(index_part) - 1  # 转换为0基索引
                
                # 解析等级和反馈
                if ' - ' in evaluation_part:
                    grade_part, feedback_part = evaluation_part.split(' - ', 1)
                    grade = grade_part.strip().lower()
                    feedback = feedback_part.strip()
                else:
                    grade = evaluation_part.strip().lower()
                    feedback = ""
                
                # 标准化等级
                if grade in ['excellent', '优秀']:
                    grade = 'excellent'
                elif grade in ['good', '良好']:
                    grade = 'good'
                elif grade in ['acceptable', '可接受']:
                    grade = 'acceptable'
                elif grade in ['poor', '较差']:
                    grade = 'poor'
                else:
                    grade = 'acceptable'  # 默认值
                
                evaluation_results.append({
                    "index": index,
                    "grade": grade,
                    "feedback": feedback
                })
                
                # 记录不合格的索引
                if grade == "poor":
                    failed_indices.append(index)
                    
            except (ValueError, IndexError):
                continue
        
        # 确保所有项目都有评估结果
        for i in range(len(state["all_subtitles"])):
            if not any(result["index"] == i for result in evaluation_results):
                evaluation_results.append({
                    "index": i,
                    "grade": "acceptable",
                    "feedback": ""
                })
        
        state["evaluation_results"] = evaluation_results
        state["failed_indices"] = failed_indices
        return state


    async def _individual_retranslator(self, state: TranslationAgentState) -> TranslationAgentState:
        """对不合格的台词进行单独重新翻译"""
        
        final_translations = state["batch_translations"].copy()
        
        for failed_index in state["failed_indices"]:
            original_subtitle = state["all_subtitles"][failed_index]
            original_text = original_subtitle.get("text", "").strip()
            feedback = state["evaluation_results"][failed_index]["feedback"]
            
            if not original_text:
                continue
                
            # 单独重新翻译
            if feedback:
                prompt = f"将以下台词翻译成{state['target_language']}，但要考虑到反馈：{feedback}\n\n原文：{original_text}"
            else:
                prompt = f"将以下台词翻译成{state['target_language']}，输出不包含翻译文本以外的内容：\n\n{original_text}"
            
            messages = [{"role": "user", "content": prompt}]
            response = await self.llm.a_run(messages)
            retranslated_text = response.content.strip()
            
            if retranslated_text:
                final_translations[failed_index] = retranslated_text
        
        state["final_translations"] = final_translations
        return state


    async def _llm_call_translator(self, state: TranslationAgentState) -> TranslationAgentState:
        """大模型翻译台词（逐句翻译模式）"""
        
        all_subtitles = state.get("all_subtitles", [])
        target_language = state.get("target_language", "")
        feedback = state.get("feedback", "")
        
        # 处理所有字幕
        for item in all_subtitles:
            original_text = item.get("text", "").strip()
            if not original_text:
                continue
                
            if feedback:
                prompt = f"将以下台词翻译成{target_language}，但要考虑到反馈：{feedback}\n\n原文：{original_text}"
            else:
                prompt = f"将以下台词翻译成{target_language}，输出不包含翻译文本以外的内容：\n\n{original_text}"
            
            messages = [{"role": "user", "content": prompt}]
            response = await self.llm.a_run(messages)
            
            translated_text = response.content.strip()
            if translated_text:
                item["text"] = translated_text
        
        return state


    async def _llm_call_evaluator(self, state: TranslationAgentState) -> TranslationAgentState:
        """LLM 评估台词翻译效果（逐句翻译模式）"""
        
        all_subtitles = state.get("all_subtitles", [])
        target_language = state.get("target_language", "")
        
        # 评估所有翻译
        quality_scores = []
        feedbacks = []
        
        for item in all_subtitles:
            original_text = item.get("text", "").strip()
            translated_text = item.get("text", "").strip()  # 这里应该是翻译后的文本
            
            if not original_text or not translated_text:
                quality_scores.append("acceptable")
                feedbacks.append("")
                continue
            
            evaluation_prompt = f"""请对译文质量打分并简述理由（尽量简短）。

请按以下JSON格式返回结果：
{{
    "grade": "excellent|good|acceptable|poor",
    "feedback": "反馈内容"
}}

原文：{original_text}
译文：{translated_text}
目标语言：{target_language}"""
            
            messages = [{"role": "user", "content": evaluation_prompt}]
            response = await self.llm.a_run(messages)
            
            # 解析评估结果
            evaluation_result = self._parse_evaluation_response(response)
            quality_scores.append(evaluation_result.get("grade", "acceptable"))
            feedbacks.append(evaluation_result.get("feedback", ""))
        
        state["quality_scores"] = quality_scores
        state["feedbacks"] = feedbacks
        return state


    def _route_translation(self, state: TranslationAgentState) -> str:
        """根据评估者的反馈返回翻译器或结束（逐句翻译模式）"""
        
        quality_scores = state.get("quality_scores", [])
        
        # 检查是否有需要重新翻译的
        for score in quality_scores:
            if score == "poor":
                return "Rejected + Feedback"
        
        return "Accepted"


    def _route_batch_processing(self, state: TranslationAgentState) -> str:
        """根据批量评估结果决定是否需要单独重新翻译"""
        
        failed_indices = state.get("failed_indices", [])
        if not failed_indices:
            return "Completed"
        else:
            return "Retranslate"


    async def _load_subtitles(self, state: TranslationAgentState) -> TranslationAgentState:
        """加载字幕文件"""
        try:
            input_path = state["input_path"]
            items = self._srt2dic(input_path)
            
            if items:
                state["all_subtitles"] = items
                print(f"成功加载字幕文件，共 {len(items)} 条字幕")
            else:
                raise ValueError("无法读取字幕内容")
        except Exception as e:
            print(f"加载字幕文件时出错: {e}")
            state["all_subtitles"] = []
        return state
    
    async def _generate_result(self, state: TranslationAgentState) -> TranslationAgentState:
        """生成最终结果"""
        try:
            all_subtitles = state.get("all_subtitles", [])
            final_translations = state.get("final_translations", [])
            failed_indices = state.get("failed_indices", [])
            
            if not all_subtitles:
                print("❌ 没有可用的字幕数据")
                state["result"] = None
                return state
            
            # 更新字幕文本
            if final_translations:
                for i, item in enumerate(all_subtitles):
                    if i < len(final_translations) and final_translations[i]:
                        item["text"] = final_translations[i]
            
            # 保存翻译结果
            output_path = state["output_path"]
            self._dic2srt(all_subtitles, output_path)
            
            state["result"] = TranslationAgentOutput(
                input_path=state["input_path"],
                output_path=output_path,
                target_language=state["target_language"],
                total_subtitles=len(all_subtitles),
                translated_subtitles=len([t for t in final_translations if t]),
                failed_translations=len(failed_indices)
            )
            
            print("✅ 字幕翻译完成")
            print(f"📊 翻译结果:")
            print(f"  📄 输入文件: {state['input_path']}")
            print(f"  📄 输出文件: {output_path}")
            print(f"  📚 总字幕数: {len(all_subtitles)}")
            print(f"  ✅ 翻译成功: {len([t for t in final_translations if t])}")
            print(f"  ❌ 翻译失败: {len(failed_indices)}")
            
        except Exception as e:
            print(f"❌ 生成结果时出错: {e}")
            state["result"] = None
        return state
    
    async def run_translation(self, input_path: str, output_path: str, target_language: str, use_batch: bool = True) -> TranslationAgentOutput:
        """运行字幕翻译"""
        initial_state = TranslationAgentState(
            input_path=input_path,
            output_path=output_path,
            target_language=target_language,
            use_batch=use_batch,
            all_subtitles=[],
            batch_translations=[],
            evaluation_results=[],
            failed_indices=[],
            final_translations=[],
            result=None
        )
        
        print(f"🚀 开始翻译字幕: {input_path}")
        
        # 选择工作流
        if use_batch:
            app = self.batch_app
        else:
            app = self.app
        
        # 运行工作流
        final_state = await app.ainvoke(initial_state)
        
        return final_state["result"]
    
    def _srt2dic(self, path: str):
        """将 SRT 文件转换为结构化数据列表（无 lines 字段），并计算 audio_duration（秒）。"""
        from utils.file_utils import read_srt as _read_srt
        blocks = _read_srt(path)
        items = []
        
        def _parse_ts(ts: str):
            try:
                if not ts:
                    return None
                hms, ms = ts.split(",")
                hours, minutes, seconds = hms.split(":")
                total = (
                    int(hours) * 3600
                    + int(minutes) * 60
                    + int(seconds)
                    + int(ms) / 1000.0
                )
                return float(total)
            except Exception:
                return None
        for block in blocks:
            if not block:
                continue
            index_value = block[0].strip()
            try:
                index_num = int(index_value)
            except Exception:
                index_num = index_value

            start_ts = ""
            end_ts = ""
            if len(block) >= 2:
                time_line = block[1]
                if "-->" in time_line:
                    parts = time_line.split("-->")
                    if len(parts) == 2:
                        start_ts = parts[0].strip()
                        end_ts = parts[1].strip()

            text_lines = block[2:] if len(block) > 2 else []
            text_joined = "\n".join([line.rstrip("\n") for line in text_lines]).strip()

            start_sec = _parse_ts(start_ts)
            end_sec = _parse_ts(end_ts)
            duration = 0.0
            if start_sec is not None and end_sec is not None:
                duration = round(max(0.0, end_sec - start_sec), 2)

            items.append({
                "index": index_num,
                "start": start_ts,
                "end": end_ts,
                "audio_duration": duration,
                "text": text_joined,
            })

        return items 
        
    def _dic2srt(self, items, path: str):
        """将结构化数据写回为 SRT 文件（仅使用 text 拆分行）。"""
        from utils.file_utils import write_srt as _write_srt
        blocks = []
        for item in items:
            index_value = item.get("index", "")
            start_ts = item.get("start", "")
            end_ts = item.get("end", "")
            text = item.get("text", "")
            lines = text.split("\n") if text else []

            time_line = f"{start_ts} --> {end_ts}".strip()
            block = [index_value, time_line]
            block.extend(lines)
            blocks.append(block)

        _write_srt(blocks, path)
    
    def _parse_evaluation_response(self, response) -> dict:
        """解析评估响应的公共函数"""
        try:
            # 提取JSON部分
            json_match = re.search(r'\{.*\}', response.content, re.DOTALL)
            if not json_match:
                print(f"❌ 评估响应：未找到JSON格式数据")
                return {"grade": "acceptable", "feedback": ""}
            
            json_str = json_match.group()
            
            # 清理字符串，去除多余的空格和换行
            cleaned_str = re.sub(r'\s+', ' ', json_str).strip()
            
            # 尝试修复可能的截断
            open_brackets = cleaned_str.count('{') - cleaned_str.count('}')
            open_squares = cleaned_str.count('[') - cleaned_str.count(']')
            cleaned_str += '}' * open_brackets
            cleaned_str += ']' * open_squares
            
            # 解析JSON
            result = json.loads(cleaned_str)
            print(f"✅ 评估响应：JSON解析成功")
            return result
            
        except json.JSONDecodeError as e:
            print(f"❌ 评估响应：JSON解析失败: {e}")
            return {"grade": "acceptable", "feedback": ""}
        except Exception as e:
            print(f"❌ 评估响应：解析过程中出错: {e}")
            return {"grade": "acceptable", "feedback": ""}
    


def translate_srt_file(input_path: str, output_path: str, target_language: str, use_batch: bool = True):
    """读取 SRT、翻译文本部分并输出新的 SRT。
    
    Args:
        input_path: 输入SRT文件路径
        output_path: 输出SRT文件路径
        target_language: 目标语言
        use_batch: 是否使用批量翻译模式，默认为True
    """
    import asyncio
    
    # 创建翻译代理配置
    config = TranslationAgentConfig(agent_name="translation_agent")
    
    # 创建翻译代理实例
    agent = TranslationAgent(config)
    
    # 运行翻译
    result = asyncio.run(agent.run_translation(input_path, output_path, target_language, use_batch))
    
    return result







